# pyrrd - Object-oriented wrapper for rrdtool.
# Copyright (C) 2008  Denis Nelubin aka Gelin
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

# http://gelin.ru
# mailto:den@gelin.ru
# xmpp:gelin@jabber.ru

""" This module provides more object-oriented wrapper for rrdtool module.
"""

import rrdtool
import math
import datetime
import time

class RRDException(Exception): pass

GAUGE = "GAUGE"
COUNTER = "COUNTER"
DERIVE = "DERIVE"
ABSOLUTE = "ABSOLUTE"
COMPUTE = "COMPUTE"

class DS:
    """ Data Set which is used to create the RRD file.

        Attributes:
        name -- string
        type -- GAUGE, COUNTER, DERIVE, ABSOLUTE or COMPUTE
        heartbeat -- seconds, usually step * 2
        min -- minimum value, default "U" - unknown
        max -- maximum value, default "U" - unknown
        rpn -- expression for the COMPUTE type
        step -- you can define step (in seconds) to autocalculate heartbeat

        The string representation of this class looks like:
        DS:name:type:heartbeat:min:max
        and so can be directly passed to the rrdtool.

        For details see 'man rrdcreate'.
    """

    def __init__(self, name, type, heartbeat=None, min=None, max=None,
            rpn=None, step=None):
        self.name = name
        self.type = type
        if heartbeat != None:
            self.heartbeat = heartbeat
        else:
            if step != None:
                self.heartbeat = step * 2
            elif self.type != COMPUTE:
                raise RRDException("'heartbeat' or 'step' must be defined")
        self.min = min
        self.max = max
        self.rpn = rpn

    def __str__(self):
        result = "DS"
        result += ":" + str(self.name)
        result += ":" + str(self.type)
        if self.type == COMPUTE:
            result += ":" + str(self.rpn)
            return result
        result += ":" + str(self.heartbeat)
        if self.min != None:
            result += ":" + str(self.min)
        else:
            result += ":U"
        if self.max != None:
            result += ":" + str(self.max)
        else:
            result += ":U"
        return result

AVERAGE = "AVERAGE"
MIN = "MIN"
MAX = "MAX"
LAST = "LAST"

DAY = 86400     #seconds
WEEK = DAY * 7
MONTH = DAY * 31
YEAR = DAY * 366

_MAX_ROWS = 1500

class RRA:
    """ Round Robin Archive definition which is used to create the RRD file.

        Attributes:
        cf -- Consolication Function name: AVERAGE, MIN, MAX or LAST
        xff -- factor to calculate *UNKNOWN* values, from 0 to 1 exclusive,
               usually 0.5
        steps -- number of primary data points which are used to build
                 a consolidated data point
        rows -- number of data points to save
        step -- you can define step (in seconds) to autocalculate steps and rows
        interval -- DAY, WEEK, MONTH, YEAR or number of seconds;
                    you can define this value to autocalculate steps and rows
        maxRows -- maximum number or rows in autocalculation, default is 1500

        Your should manually define 'steps' and 'rows' values or use
        autocalculation defining 'step' and 'interval'
        (and optionally 'maxRows') value.

        The string representation of this class looks like:
        RRA:cf:xff:steps:rows
        and so can be directly passed to the rrdtool.

        For details see 'man rrdcreate'.
    """

    def __init__(self, cf, xff=0.5, steps=None, rows=None,
            step=None, interval=None, maxRows=_MAX_ROWS):
        self.cf = cf
        if xff != None:
            self.xff = xff
        else:
            self.xff = 0.5
        if steps != None:
            self.steps = steps
        else:
            self.steps = self._defineSteps(step, interval, maxRows)
        if rows != None:
            self.rows = rows
        else:
            self.rows = self._defineRows(step, interval)

    def __str__(self):
        result = "RRA"
        result += ":" + str(self.cf)
        result += ":" + str(self.xff)
        result += ":" + str(self.steps)
        result += ":" + str(self.rows)
        return result

    def _defineStepsAndRows(self, step, interval, maxRows):
        """ Finds the suitable number of PDP to built CDP,
            and the number of CDP rows to represent
            the specified data interval.
            interval / (steps * step) = rows <= 1500
        """
        if step == None or interval == None:
            raise RRDException(
                    "both 'steps' and 'rows' or both 'step' and 'interval'" +
                    " must be defined")
        # steps >= interval / (step * 1500)
        steps = int(math.ceil(float(interval) / (step * maxRows)))
        rows = int(interval / (steps * step))
        return (steps, rows)

    def _defineSteps(self, step, interval, maxRows):
        """ Finds the suitable number of PDP to built CDP
            to represent the specified data interval.
            interval / (steps * step) = rows <= 1500
        """
        if step == None or interval == None:
            raise RRDException(
                    "'steps' or both 'step' and 'interval'" +
                    " must be defined")
        # steps >= interval / (step * 1500)
        return int(math.ceil(float(interval) / (step * maxRows)))

    def _defineRows(self, step, interval):
        """ Finds the suitable number of CDP rows to represent
            the specified data interval.
            This method uses predefined steps value.
            interval / (steps * step) = rows <= 1500
        """
        if step == None or interval == None:
            raise RRDException(
                    "both 'steps' and 'rows' or both 'step' and 'interval'" +
                    " must be defined")
        # steps >= interval / (step * 1500)
        return int(interval / (self.steps * step))

class Timestamp(datetime.datetime):
    """ Simple extension of the datetime class, displays timestamp
        (seconds from epoch) when converting to string
    """

    def __str__(self):
        return str(int(time.mktime(self.timetuple())))

class RRD:
    """ Represents RRD database file.
    """

    def __init__(self, filename):
        """ Creates RRD object

            Arguments:
            filename -- Location of the RRD file.
        """
        self.filename = filename

    def update(self, *args, **kwargs):
        """ Updates RRD database.

            Arguments:
            timestamp -- Timestamp object or number of seconds since the Epoch
                or just a string in 'at style' (see man rrdfetch),
                default value is 'N' which means 'now'
            All other positional or keyword arguments (but not both!)
            represents values to be inserted into the database.
            Order of the positional arguments or the names of
            the keyword arguments must be the same as
            were defined in RRA when the database file was created.
        """
        (values, template) = self._parseUpdateParams(args, kwargs)
        if template != '':
            rrdtool.update(self.filename, '-t', template, values)
        else:
            rrdtool.update(self.filename, values)

    def _parseUpdateParams(self, args, kwargs):
        """ Creates values and, optionally, template -- parameters
            for the rrdtool.update -- from arguments to update
            method.
        """
        timestamp = 'N'
        keyargs = kwargs
        if keyargs.has_key('timestamp'):
            timestamp = keyargs.get('timestamp')
            del keyargs['timestamp']
        result = str(timestamp)
        if _isAtString(timestamp):
            result += '@'
        else:
            result += ':'

        values = []
        for data in keyargs.values():
            if data == None:
                values.append('U')
            else:
                values.append(str(data))
        for data in args:
            if data == None:
                values.append('U')
            else:
                values.append(str(data))

        result += ':'.join(values)
        template = ':'.join(kwargs.keys())

        return (result, template)

def create(filename, start=None, step=None, ds=[], rra=[]):
    """ Creates new RRD file with specified parameters.

        Arguments:
        filename -- name of the RRD file
        start -- time of the last update for new file, can be Timestamp or
                 number of seconds sience epoch or any "at string"
                 (see 'man rrdfetch')
        step -- number of seconds between data updates
        ds -- list of Data Set definitions, each definition can be a DS object,
              or just a string in rrdtool-specific syntax
        rra -- list of Round Robin Archive definitions, each definition can be
               an RRA object or just a string in rrdtool-specific syntax
        For details see 'man rrdcreate'.

        Returns RRD object, ready to update or fetch.
    """
    args = []
    args.append(str(filename))
    if start != None:
        args.append("--start")
        args.append(str(start))
    if step != None:
        args.append("--step")
        args.append(str(step))
    for dsItem in ds:
        args.append(str(dsItem))
    for rraItem in rra:
        args.append(str(rraItem))
    rrdtool.create(*args)
    return RRD(filename)

def open(filename):
    """ Opens existed RRD file for the following manipulations.
        The same as RRD(filename).

        Arguments:
        filename -- name of the RRD file

        Returns RRD object.
    """
    return RRD(filename)

def _isAtString(timestamp):
    """ Checks is the timestamp "at string" or not.
        Symply supposes that "N" or number is a not "at string".
    """
    if timestamp == "N":
        return False
    try:
        int(timestamp)
    except:
        return True
    return False